"""
RAG核心引擎模块

这个模块负责：
1. 整合检索和生成功能
2. 问答流程控制
3. 结果后处理和优化

作者：xiaodeme@163.com
"""

import logging
import time
from typing import List, Dict, Any, Optional
from .config import config
from .document_processor import PDFProcessor, TextCleaner, TextSplitter
from .embeddings import EmbeddingGenerator, VectorStore
from .retriever import RetrievalEngine
from .generator import LLMClient, PromptTemplate, ResponseGenerator

# 配置日志
logging.basicConfig(level=getattr(logging, config.LOG_LEVEL))
logger = logging.getLogger(__name__)

class RAGEngine:
    """
    RAG引擎主类
    
    负责：
    - 整合文档处理、向量检索、回答生成
    - 管理完整的RAG流程
    - 提供统一的问答接口
    """
    
    def __init__(self):
        """初始化RAG引擎"""
        logger.info("正在初始化RAG引擎...")
        
        # 初始化各个组件
        self.pdf_processor = None
        self.text_cleaner = TextCleaner()
        self.text_splitter = TextSplitter()
        self.embedding_generator = EmbeddingGenerator()
        self.retrieval_engine = RetrievalEngine()
        self.llm_client = LLMClient()
        self.prompt_template = PromptTemplate()
        self.response_generator = ResponseGenerator()
        
        # 状态管理
        self.is_initialized = False
        self.document_processed = False
        
        logger.info("✅ RAG引擎初始化完成")
    
    def initialize_from_pdf(self, pdf_path: str = None) -> Dict[str, Any]:
        """
        从PDF文件初始化RAG引擎
        
        Args:
            pdf_path (str, optional): PDF文件路径，默认使用配置中的路径
            
        Returns:
            Dict[str, Any]: 初始化结果
        """
        logger.info("开始从PDF文件初始化RAG引擎...")
        
        try:
            start_time = time.time()
            
            # 1. 初始化PDF处理器
            self.pdf_processor = PDFProcessor(pdf_path)
            logger.info("✅ PDF处理器初始化完成")
            
            # 2. 提取PDF文本
            logger.info("正在提取PDF文本...")
            raw_text = self.pdf_processor.extract_text()
            
            # 3. 清洗文本
            logger.info("正在清洗文本...")
            cleaned_text = self.text_cleaner.clean_text(raw_text)
            
            # 4. 分割文本
            logger.info("正在分割文本...")
            text_chunks = self.text_splitter.split_text(cleaned_text)
            
            # 5. 添加文档到检索引擎
            logger.info(f"正在处理 {len(text_chunks)} 个文本块...")
            metadata = [{'chunk_id': i, 'source': 'pdf'} for i in range(len(text_chunks))]
            self.retrieval_engine.add_documents(text_chunks, metadata)
            
            # 6. 测试LLM连接
            logger.info("测试LLM连接...")
            connection_result = self.llm_client.test_connection()
            
            processing_time = time.time() - start_time
            
            result = {
                'success': True,
                'message': 'RAG引擎初始化成功',
                'statistics': {
                    'original_text_length': len(raw_text),
                    'cleaned_text_length': len(cleaned_text),
                    'text_chunks_count': len(text_chunks),
                    'processing_time': processing_time,
                    'llm_available': connection_result['success']
                },
                'pdf_info': self.pdf_processor.get_pdf_info(),
                'retrieval_stats': self.retrieval_engine.get_statistics()
            }
            
            self.is_initialized = True
            self.document_processed = True
            
            logger.info(f"✅ RAG引擎初始化成功，耗时: {processing_time:.2f}秒")
            return result
            
        except Exception as e:
            logger.error(f"❌ RAG引擎初始化失败: {str(e)}")
            return {
                'success': False,
                'message': f'初始化失败: {str(e)}',
                'error': str(e)
            }
    
    def ask_question(self, question: str, 
                    use_retrieval: bool = True,
                    max_context_length: int = 2000) -> Dict[str, Any]:
        """
        回答用户问题
        
        Args:
            question (str): 用户问题
            use_retrieval (bool): 是否使用检索增强
            max_context_length (int): 最大上下文长度
            
        Returns:
            Dict[str, Any]: 回答结果
        """
        if not question or not question.strip():
            return {
                'success': False,
                'message': '问题不能为空',
                'question': question,
                'answer': ''
            }
        
        logger.info(f"收到问题: '{question[:100]}...'")
        
        try:
            start_time = time.time()
            
            if use_retrieval and self.document_processed:
                # RAG模式：检索相关文档后生成回答
                result = self._answer_with_retrieval(question, max_context_length)
            else:
                # 简单模式：直接使用LLM生成回答
                result = self._answer_without_retrieval(question)
            
            processing_time = time.time() - start_time
            result['processing_time'] = processing_time
            
            logger.info(f"问答完成，耗时: {processing_time:.2f}秒")
            return result
            
        except Exception as e:
            logger.error(f"❌ 问答处理失败: {str(e)}")
            return {
                'success': False,
                'message': f'问答处理失败: {str(e)}',
                'question': question,
                'answer': '',
                'error': str(e)
            }
    
    def _answer_with_retrieval(self, question: str, max_context_length: int) -> Dict[str, Any]:
        """
        使用检索增强生成回答
        
        Args:
            question (str): 用户问题
            max_context_length (int): 最大上下文长度
            
        Returns:
            Dict[str, Any]: 回答结果
        """
        logger.debug("使用RAG模式生成回答")
        
        # 1. 检索相关文档
        retrieval_results = self.retrieval_engine.search_and_rank(
            question, 
            top_k=config.TOP_K,
            similarity_threshold=config.SIMILARITY_THRESHOLD
        )
        
        if not retrieval_results:
            logger.warning("未找到相关文档，使用简单模式")
            return self._answer_without_retrieval(question)
        
        # 2. 构造上下文
        context_parts = []
        total_length = 0
        
        for result in retrieval_results:
            doc_text = result.get('content', result.get('document', ''))
            if total_length + len(doc_text) <= max_context_length:
                context_parts.append(doc_text)
                total_length += len(doc_text)
            else:
                # 截断最后一个文档以适应长度限制
                remaining_length = max_context_length - total_length
                if remaining_length > 50:  # 确保有足够的空间
                    context_parts.append(doc_text[:remaining_length])
                break
        
        context = "\n\n".join(context_parts)
        
        # 3. 生成回答
        rag_result = self.response_generator.generate_rag_answer(question, context)
        
        # 4. 构造完整结果
        result = {
            'success': rag_result['success'],
            'message': 'RAG问答完成' if rag_result['success'] else rag_result.get('error', '生成失败'),
            'question': question,
            'answer': rag_result['response'],
            'mode': 'rag',
            'retrieval_results': retrieval_results,
            'context_used': context,
            'context_length': len(context),
            'llm_info': rag_result.get('llm_info', {})
        }
        
        return result
    
    def _answer_without_retrieval(self, question: str) -> Dict[str, Any]:
        """
        不使用检索，直接生成回答
        
        Args:
            question (str): 用户问题
            
        Returns:
            Dict[str, Any]: 回答结果
        """
        logger.debug("使用简单模式生成回答")
        
        # 直接使用LLM生成回答
        simple_result = self.response_generator.generate_simple_answer(question)
        
        # 构造结果
        result = {
            'success': simple_result['success'],
            'message': '简单问答完成' if simple_result['success'] else simple_result.get('error', '生成失败'),
            'question': question,
            'answer': simple_result['response'],
            'mode': 'simple',
            'retrieval_results': [],
            'context_used': '',
            'context_length': 0,
            'llm_info': simple_result.get('llm_info', {})
        }
        
        return result
    
    def get_system_status(self) -> Dict[str, Any]:
        """
        获取系统状态信息
        
        Returns:
            Dict[str, Any]: 系统状态
        """
        try:
            # 测试LLM连接
            llm_status = self.llm_client.test_connection()
            
            # 获取检索引擎统计
            retrieval_stats = self.retrieval_engine.get_statistics()
            
            status = {
                'rag_engine_initialized': self.is_initialized,
                'document_processed': self.document_processed,
                'llm_available': llm_status['success'],
                'llm_message': llm_status['message'],
                'retrieval_engine': retrieval_stats,
                'configuration': {
                    'chunk_size': config.CHUNK_SIZE,
                    'chunk_overlap': config.CHUNK_OVERLAP,
                    'top_k': config.TOP_K,
                    'similarity_threshold': config.SIMILARITY_THRESHOLD,
                    'ollama_model': config.OLLAMA_MODEL
                }
            }
            
            if self.pdf_processor:
                status['pdf_info'] = self.pdf_processor.get_pdf_info()
            
            return status
            
        except Exception as e:
            logger.error(f"获取系统状态失败: {str(e)}")
            return {
                'error': str(e),
                'rag_engine_initialized': False,
                'document_processed': False,
                'llm_available': False
            }
    
    def clear_and_reset(self):
        """清空并重置RAG引擎"""
        logger.info("正在清空并重置RAG引擎...")
        
        try:
            # 清空检索引擎
            self.retrieval_engine.clear_all()
            
            # 重置状态
            self.is_initialized = False
            self.document_processed = False
            self.pdf_processor = None
            
            logger.info("✅ RAG引擎已重置")
            
        except Exception as e:
            logger.error(f"❌ 重置失败: {str(e)}")

class QAProcessor:
    """
    问答处理器类
    
    提供更高级的问答处理功能
    """
    
    def __init__(self, rag_engine: RAGEngine):
        """
        初始化问答处理器
        
        Args:
            rag_engine (RAGEngine): RAG引擎实例
        """
        self.rag_engine = rag_engine
        self.conversation_history = []
        
        logger.info("初始化问答处理器")
    
    def process_question(self, question: str, 
                        remember_conversation: bool = True) -> Dict[str, Any]:
        """
        处理问题并记录对话历史
        
        Args:
            question (str): 用户问题
            remember_conversation (bool): 是否记录对话历史
            
        Returns:
            Dict[str, Any]: 处理结果
        """
        logger.info(f"处理问题: '{question[:50]}...'")
        
        try:
            # 处理问题
            result = self.rag_engine.ask_question(question)
            
            # 记录对话历史
            if remember_conversation:
                conversation_entry = {
                    'timestamp': time.time(),
                    'question': question,
                    'answer': result.get('answer', ''),
                    'success': result.get('success', False),
                    'mode': result.get('mode', 'unknown')
                }
                self.conversation_history.append(conversation_entry)
                
                # 限制历史记录数量
                if len(self.conversation_history) > 50:
                    self.conversation_history = self.conversation_history[-50:]
            
            # 添加对话历史信息到结果中
            result['conversation_id'] = len(self.conversation_history)
            
            return result
            
        except Exception as e:
            logger.error(f"问题处理失败: {str(e)}")
            return {
                'success': False,
                'question': question,
                'answer': '',
                'error': str(e)
            }
    
    def get_conversation_history(self, limit: int = 10) -> List[Dict[str, Any]]:
        """
        获取对话历史
        
        Args:
            limit (int): 返回的对话数量限制
            
        Returns:
            List[Dict[str, Any]]: 对话历史列表
        """
        return self.conversation_history[-limit:] if self.conversation_history else []
    
    def clear_conversation_history(self):
        """清空对话历史"""
        self.conversation_history.clear()
        logger.info("对话历史已清空")

def test_rag_engine():
    """
    测试RAG引擎
    """
    print("=== 测试RAG引擎 ===")
    
    try:
        # 创建RAG引擎
        rag_engine = RAGEngine()
        print("✅ RAG引擎创建成功")
        
        # 测试系统状态
        status = rag_engine.get_system_status()
        print(f"✅ 系统状态获取成功")
        print(f"   LLM可用: {status.get('llm_available', False)}")
        print(f"   文档已处理: {status.get('document_processed', False)}")
        
        # 如果有PDF文件，尝试初始化
        try:
            init_result = rag_engine.initialize_from_pdf()
            if init_result['success']:
                print("✅ 从PDF初始化成功")
                print(f"   文本块数量: {init_result['statistics']['text_chunks_count']}")
                
                # 测试问答
                test_question = "翁氏家族有什么特点？"
                qa_result = rag_engine.ask_question(test_question)
                
                if qa_result['success']:
                    print("✅ 问答测试成功")
                    print(f"   问题: {qa_result['question']}")
                    print(f"   回答: {qa_result['answer'][:100]}...")
                    print(f"   模式: {qa_result['mode']}")
                else:
                    print(f"⚠️ 问答测试失败: {qa_result.get('message', '未知错误')}")
            else:
                print(f"⚠️ PDF初始化失败: {init_result.get('message', '未知错误')}")
                
        except Exception as e:
            print(f"⚠️ PDF处理失败 (这是正常的，如果PDF文件不存在): {str(e)}")
        
        # 测试问答处理器
        qa_processor = QAProcessor(rag_engine)
        simple_result = qa_processor.process_question("你好，请介绍一下自己")
        
        if simple_result['success']:
            print("✅ 问答处理器测试成功")
        else:
            print(f"⚠️ 问答处理器测试失败: {simple_result.get('message', '未知错误')}")
        
        return True
        
    except Exception as e:
        print(f"❌ 测试失败: {str(e)}")
        return False

if __name__ == "__main__":
    test_rag_engine() 